/**
 * @file main.cpp
 * @brief Main file for the project
 */

#include <iostream>
#include <fstream>
#include <vector>
#include <random>
#include <math.h>

#include <dirent.h>
#include <cassert>
#include <sys/time.h>

#include "acl/acl.h"

#define MATRIX_M 256
#define MATRIX_K 32
#define MATRIX_N 128
#define ROW_SHAPE 1
#define COL_SHAPE 8
#define THOUSAND 1000
#define TIMES_UNIT 1000000
#define MULTIPLY_ADD 2             // mmad 矩阵乘加，两次浮点运算
#define TRANSFER_TFLOPS 1000000000 // FLOPS换算TFLOPS

// 910A core_num 32
int aiCoreCount = 32;

int executeTimes = 100;

#define OP_FILE "train_mmad_model.o"
#define OP_KERNEL_NAME "train_mmad_model"
#define OP_KERNEL_ID "train_mmad_model__kernel0"

struct OperatorParam
{
  int32_t m;
  int32_t k;
  int32_t n;
  int32_t ratio;
  int32_t duration;
};

#define ACL_REQUIRES_OK(expr)            \
  do                                     \
  {                                      \
    auto __ret = (expr);                 \
    if (__ret != ACL_SUCCESS)            \
    {                                    \
      printf("error at %d\n", __LINE__); \
      return __ret;                      \
    }                                    \
  } while (0)

bool CheckDims(size_t n, size_t c, size_t h, size_t w)
{
  if (n == 0 || c == 0 || h == 0 || w == 0)
  {
    std::cout << "[ERROR] All dim can't support 0" << std::endl;
    return false;
  }

  if (n != 1)
  {
    std::cout << "[ERROR] Batch N only support 1" << std::endl;
    return false;
  }

  return true;
}

void DisplayResult(int64_t &n, int64_t &c, int64_t &h, int64_t &w, aclFloat16 *output)
{
  std::cout << "The result of BatchNorm is(showing no more than 16 elements):" << std::endl;
  int i = 0;
  for (; i < (n * c * h * w); ++i)
  {
    if (i < 16)
    {
      std::cout << aclFloat16ToFloat(output[i]) << ", ";
    }
    if (aclFloat16ToFloat(output[i]) != 5.0)
    {
      printf("Case failed! Index %d element is %f but 5.0 expected.\n",
             i, aclFloat16ToFloat(output[i]));
      break;
    }
  }
  if (i == (n * c * h * w))
  {
    std::cout << "Case passed!" << std::endl;
  }
}

aclError run_test()
{
  printf("start run test\n");
  aclrtRunMode runMode;
  if (aclrtGetRunMode(&runMode) != ACL_SUCCESS)
  {
    std::cout << "[ERROR] Get run mode failed" << std::endl;
    return ACL_ERROR_BAD_ALLOC;
  }

  aclTensorDesc *input_desc[4];
  aclTensorDesc *output_desc[1];

  std::vector<int64_t> shapeA = {MATRIX_M, MATRIX_K};
  std::vector<int64_t> shapeB = {MATRIX_K, MATRIX_N};
  std::vector<int64_t> shapeC = {MATRIX_M, MATRIX_N};
  std::vector<int64_t> shapeRatio = {ROW_SHAPE, COL_SHAPE};
  std::vector<int64_t> shapeDuration = {ROW_SHAPE, COL_SHAPE};

  input_desc[0] = aclCreateTensorDesc(ACL_FLOAT16, shapeA.size(), shapeA.data(), ACL_FORMAT_ND);
  input_desc[1] = aclCreateTensorDesc(ACL_FLOAT16, shapeB.size(), shapeB.data(), ACL_FORMAT_ND);
  input_desc[2] = aclCreateTensorDesc(ACL_INT32, shapeRatio.size(), shapeRatio.data(), ACL_FORMAT_ND);
  input_desc[3] = aclCreateTensorDesc(ACL_INT32, shapeDuration.size(), shapeDuration.data(), ACL_FORMAT_ND);
  output_desc[0] = aclCreateTensorDesc(ACL_FLOAT, shapeC.size(), shapeC.data(), ACL_FORMAT_ND);

  aclopAttr *attr = aclopCreateAttr();
  aclopSetAttrInt(attr, "type", 1);

  aclrtStream stream;
  ACL_REQUIRES_OK(aclrtCreateStream(&stream));

  void *InputA = nullptr;
  void *InputB = nullptr;
  void *InputRatio = nullptr;
  void *InputDuration = nullptr;
  void *Output = nullptr;

  ACL_REQUIRES_OK(aclrtMalloc(&InputA, aclGetTensorDescSize(input_desc[0]), ACL_MEM_MALLOC_NORMAL_ONLY));
  ACL_REQUIRES_OK(aclrtMalloc(&InputB, aclGetTensorDescSize(input_desc[1]), ACL_MEM_MALLOC_NORMAL_ONLY));
  ACL_REQUIRES_OK(aclrtMalloc(&InputRatio, aclGetTensorDescSize(input_desc[2]), ACL_MEM_MALLOC_NORMAL_ONLY));
  ACL_REQUIRES_OK(aclrtMalloc(&InputDuration, aclGetTensorDescSize(input_desc[3]), ACL_MEM_MALLOC_NORMAL_ONLY));
  ACL_REQUIRES_OK(aclrtMalloc(&Output, aclGetTensorDescSize(output_desc[0]), ACL_MEM_MALLOC_NORMAL_ONLY));

  aclrtMemcpyKind kindInput = ACL_MEMCPY_HOST_TO_DEVICE;
  if (runMode == ACL_DEVICE)
  {
    kindInput = ACL_MEMCPY_DEVICE_TO_DEVICE;
  }

  size_t InputA_Size = aclGetTensorDescSize(input_desc[0]);
  size_t InputB_Size = aclGetTensorDescSize(input_desc[1]);
  size_t InputR_Size = aclGetTensorDescSize(input_desc[2]);
  size_t InputD_Size = aclGetTensorDescSize(input_desc[3]);

  aclFloat16 *xxxA = (aclFloat16 *)malloc(aclGetTensorDescSize(input_desc[0]));
  aclFloat16 *xxxB = (aclFloat16 *)malloc(aclGetTensorDescSize(input_desc[1]));
  int32_t *xxxR = (int32_t *)malloc(aclGetTensorDescSize(input_desc[2]));
  int32_t *xxxD = (int32_t *)malloc(aclGetTensorDescSize(input_desc[3]));

  for (int i = 0; i < aclGetTensorDescSize(input_desc[0]) / sizeof(ACL_FLOAT16); ++i)
  {
    xxxA[i] = aclFloatToFloat16((random() % 100) * M_PI);
  }
  for (int i = 0; i < aclGetTensorDescSize(input_desc[1]) / sizeof(ACL_FLOAT16); ++i)
  {
    xxxB[i] = aclFloatToFloat16((random() % 100) * M_PI);
  }
  for (int i = 0; i < aclGetTensorDescSize(input_desc[2]) / sizeof(ACL_INT32); ++i)
  {
    xxxR[i] = 0;
  }

  for (int i = 0; i < aclGetTensorDescSize(input_desc[3]) / sizeof(ACL_INT32); ++i)
  {
    xxxD[i] = executeTimes * TIMES_UNIT;
  }

  aclrtMemcpy(InputA, InputA_Size, xxxA, InputA_Size, kindInput);
  aclrtMemcpy(InputB, InputB_Size, xxxB, InputB_Size, kindInput);
  aclrtMemcpy(InputRatio, InputR_Size, xxxR, InputR_Size, kindInput);
  aclrtMemcpy(InputDuration, InputD_Size, xxxD, InputD_Size, kindInput);

  aclDataBuffer *bufA = aclCreateDataBuffer(InputA, aclGetTensorDescSize(input_desc[0]));
  aclDataBuffer *bufB = aclCreateDataBuffer(InputB, aclGetTensorDescSize(input_desc[1]));
  aclDataBuffer *bufRatio = aclCreateDataBuffer(InputRatio, aclGetTensorDescSize(input_desc[2]));
  aclDataBuffer *bufDuration = aclCreateDataBuffer(InputDuration, aclGetTensorDescSize(input_desc[3]));
  aclDataBuffer *outputBuffer = aclCreateDataBuffer(Output, aclGetTensorDescSize(output_desc[0]));

  aclDataBuffer *inputs[] = {bufA, bufB, bufRatio, bufDuration};
  aclDataBuffer *outputs[] = {outputBuffer};

  auto ret = aclopUpdateParams(OP_KERNEL_NAME, 4, input_desc, 1, output_desc, attr);
  if (ret != ACL_ERROR_NONE)
  {
    printf("aclopUpdateParams error:%d\n", ret);
    return ret;
  }

  struct timeval begin = {0};
  struct timeval end = {0};
  gettimeofday(&begin, NULL);
  ACL_REQUIRES_OK(aclopExecuteV2(OP_KERNEL_NAME,
                                 4, input_desc, inputs,
                                 1, output_desc, outputs,
                                 attr, stream));
  ACL_REQUIRES_OK(aclrtSynchronizeStream(stream));
  gettimeofday(&end, NULL);

  int escaped = ((end.tv_sec * THOUSAND) + (end.tv_usec / THOUSAND)) - ((begin.tv_sec * THOUSAND) + (begin.tv_usec / THOUSAND));
  float work_load = static_cast<float>(aiCoreCount) * MATRIX_M * MATRIX_K * MATRIX_N * executeTimes * TIMES_UNIT * MULTIPLY_ADD;
  float flops = work_load / escaped / TRANSFER_TFLOPS;

  printf("Used Time = %d Flops = %f T\n", escaped, flops);

  ACL_REQUIRES_OK(aclrtDestroyStream(stream));
  aclrtMemcpyKind kindOutput = ACL_MEMCPY_DEVICE_TO_HOST;
  if (runMode == ACL_DEVICE)
  {
    kindOutput = ACL_MEMCPY_DEVICE_TO_DEVICE;
  }
  ACL_REQUIRES_OK(aclrtFree(InputA));
  ACL_REQUIRES_OK(aclrtFree(InputB));
  ACL_REQUIRES_OK(aclrtFree(InputRatio));
  ACL_REQUIRES_OK(aclrtFree(InputDuration));
  ACL_REQUIRES_OK(aclrtFree(Output));

  ACL_REQUIRES_OK(aclDestroyDataBuffer(bufA));
  ACL_REQUIRES_OK(aclDestroyDataBuffer(bufB));
  ACL_REQUIRES_OK(aclDestroyDataBuffer(bufRatio));
  ACL_REQUIRES_OK(aclDestroyDataBuffer(bufDuration));
  ACL_REQUIRES_OK(aclDestroyDataBuffer(outputBuffer));

  aclDestroyTensorDesc(input_desc[0]);
  aclDestroyTensorDesc(input_desc[1]);
  aclDestroyTensorDesc(input_desc[2]);
  aclDestroyTensorDesc(input_desc[3]);
  aclDestroyTensorDesc(output_desc[0]);

  free(xxxA);
  free(xxxB);
  free(xxxR);
  free(xxxD);

  return ACL_SUCCESS;
}

bool ReadBytesFromBinaryFile(const char *file_name, char **buffer, int &length)
{
  if (file_name == nullptr)
  {
    std::cout << "File is nullptr, parameter error." << std::endl;
    return false;
  }

  std::string real_path = file_name;
  std::ifstream file(real_path.c_str(), std::ios::binary | std::ios::ate);
  if (!file.is_open())
  {
    printf("Read file %s failed.", file_name);
    return false;
  }

  length = static_cast<int>(file.tellg());
  if (length <= 0)
  {
    printf("File length <= 0");
    file.close();
    return false;
  }

  file.seekg(0, std::ios::beg);

  *buffer = new (std::nothrow) char[length]();
  if (buffer == nullptr)
  {
    printf("New an object failed");
    file.close();
    return false;
  }

  file.read(*buffer, length);
  file.close();
  return true;
}

void Deallocator(void *data, size_t size)
{
  auto addr = reinterpret_cast<char *>(data);
  delete[] addr;
}

int main(int argc, char **argv)
{
  if (aclInit(nullptr) != ACL_SUCCESS)
  {
    std::cout << "Acl init failed." << std::endl;
    return -1;
  }

  aclrtSetDevice(0);
  aclopRegisterCompileFunc(OP_KERNEL_NAME, [](int numInputs, const aclTensorDesc *const inputDesc[], int32_t numOutputs, const aclTensorDesc *const outputDesc[], const aclopAttr *opAttr, aclopKernelDesc *aclopKernelDesc) -> aclError
                           {
		    OperatorParam param = {
		    	.m = MATRIX_M,
		    	.k = MATRIX_K,
		    	.n = MATRIX_N,
			.ratio = 0,
			.duration = 0
		    };
		    auto ret = aclopSetKernelArgs(aclopKernelDesc, 
				    OP_KERNEL_ID,
				    aiCoreCount, &param, sizeof(OperatorParam));
		    return ret; });

  const char *op_file = OP_FILE;
  char *buffer = nullptr;
  int length = 0;
  ReadBytesFromBinaryFile(op_file, &buffer, length);
  auto ret = aclopCreateKernel(OP_KERNEL_NAME,
                               OP_KERNEL_ID,
                               OP_KERNEL_ID,
                               buffer, length, ACL_ENGINE_AICORE, Deallocator);
  if (ret != ACL_SUCCESS)
  {
    printf("Create kernel failed %d\n", ret);
  }
  else
  {
    printf("Create kernel success\n");
  }

  if (run_test() != ACL_SUCCESS)
  {
    std::cout << "[ERROR][" << __LINE__ << "] Test failed." << std::endl;
    return -1;
  }

  aclrtResetDevice(0);

  if (aclFinalize() != ACL_SUCCESS)
  {
    std::cout << "[ERROR] Finalize acl failed." << std::endl;
    return -1;
  }
  return 0;
}
